package com.cy.embeddingdemo;

import dev.langchain4j.community.model.zhipu.ZhipuAiImageModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import io.milvus.client.MilvusClient;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import jakarta.annotation.Resource;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;

import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

@SpringBootTest
public class testchat {

    @Test
    public void test() {
        OpenAiChatModel model = OpenAiChatModel.builder().baseUrl("https://api.deepseek.com").apiKey("sk-dc9cbd9f379f470c9193617869600956").modelName("deepseek-chat").build();
        String chat = model.chat("你好，你是谁");
        System.out.println(chat);


    }


    // 调用本得的 ollama 服务
    @Test
    public void testOllama() {

        ChatLanguageModel model = OllamaChatModel.builder()
                .baseUrl("http://192.168.0.105:11434/")
                .modelName("qwen2.5:1.5b")
                .build();
        String answer = model.chat("List top 10 cites in China");
        System.out.println(answer);


        String json = model.chat("List top 10 cites in US");
        System.out.println(json);

    }

    //测试文生图 api
    @Test
    public void testTextToPicture() {
        ZhipuAiImageModel zhipuModel = ZhipuAiImageModel.builder()
                .apiKey("7b5ac41a489b9c31664cf218336741b1.3Hqea5I3Nf7z5s18")
                .model("cogview-3-flash")
                .callTimeout(Duration.ofSeconds(100))
                .connectTimeout(Duration.ofSeconds(100))
                .writeTimeout(Duration.ofSeconds(100))
                .readTimeout(Duration.ofSeconds(100))
                .build();
        Response<Image> imageResponse = zhipuModel.generate("生成一个美女照片");

        // 假设 Image 对象有一个方法可以获取图像的字节数组
        byte[] imageBytes = imageResponse.content().base64Data().getBytes();


        // 定义保存文件的路径
        String filePath = "./image.jpg";

        try (OutputStream outputStream = new FileOutputStream(filePath)) {
            // 将图像字节数组写入文件
            outputStream.write(imageBytes);
            System.out.println("图像已成功保存到: " + filePath);
        } catch (IOException e) {
            System.err.println("保存图像时出现错误: " + e.getMessage());
        }

    }

    @Resource
    public OpenAiEmbeddingModel embeddingModel;

    @Test
    public void test02()  {
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();




        // 利用向量模型进行向量化， 然后存储向量到向量数据库
        TextSegment segment1 = TextSegment.from("""
                预订航班:
                - 通过我们的网站或移动应用程序预订。
                - 预订时需要全额付款。
                - 确保个人信息（姓名、ID 等）的准确性，因为更正可能会产生 25 的费用。
                """);
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        embeddingStore.add(embedding1, segment1);

        // 利用向量模型进行向量化， 然后存储向量到向量数据库
        TextSegment segment2 = TextSegment.from("""
                取消预订:
                - 最晚在航班起飞前 48 小时取消。
                - 取消费用：经济舱 75 美元，豪华经济舱 50 美元，商务舱 25 美元。
                - 退款将在 7 个工作日内处理。
                """);
        Embedding embedding2 = embeddingModel.embed(segment2).content();
        embeddingStore.add(embedding2, segment2);


        // 需要查询的内容 向量化
        Embedding queryEmbedding = embeddingModel.embed("退票要多少钱").content();

        // 去向量数据库查询
        // 构建查询条件
        EmbeddingSearchRequest build = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                // .minScore(0.8)
                .build();

        // 查询
        EmbeddingSearchResult<TextSegment> segmentEmbeddingSearchResult = embeddingStore.search(build);
        segmentEmbeddingSearchResult.matches().forEach(embeddingMatch -> {
            System.out.println(embeddingMatch.score()); // 0.8144288515898701
            System.out.println(embeddingMatch.embedded().text()); // I like football.
        });

    }





}
